// modbus_tcp_server.c
#include <sys/un.h>
#include <stdio.h>  
#include <stdlib.h>  
#include <string.h>  
#include <unistd.h>  
#include <arpa/inet.h>  
#include <sys/socket.h>  
#include <pthread.h>  
#include <signal.h>  

#define PORT 502
#define MAX_CLIENTS 100
#define COIL_COUNT 1000
#define DISCRETE_INPUT_COUNT 1000
#define HOLDING_REGISTER_COUNT 1000
#define INPUT_REGISTER_COUNT 1000

#define SOCKET_PATH "/tmp/modbus_socket"  

// 共享数据定义
uint8_t coils[COIL_COUNT] = {0};
uint8_t discrete_inputs[DISCRETE_INPUT_COUNT] = {0};
uint16_t holding_registers[HOLDING_REGISTER_COUNT] = {0};
uint16_t input_registers[INPUT_REGISTER_COUNT] = {0};

// 互斥锁
pthread_mutex_t data_mutex = PTHREAD_MUTEX_INITIALIZER;

void sig_handler(int signo) {
    if (signo == SIGINT) {
        printf("Received SIGINT, exiting safely...\n");
        exit(0);
    }
}

void* handle_client(void *arg) {
    // int client_socket = (intptr_t)arg;
    int client_socket = (intptr_t)arg;  
    struct sockaddr_in address;  
    socklen_t addr_len = sizeof(address);  
    getpeername(client_socket, (struct sockaddr *)&address, &addr_len);  
    
    char client_ip[INET_ADDRSTRLEN];  
    inet_ntop(AF_INET, &address.sin_addr, client_ip, sizeof(client_ip));  

    // 发送客户端地址信息到Unix域套接字  
    int unix_socket = socket(AF_UNIX, SOCK_STREAM, 0);  
    struct sockaddr_un addr;  
    addr.sun_family = AF_UNIX;  
    strncpy(addr.sun_path, SOCKET_PATH, sizeof(addr.sun_path) - 1);  
    
    // 这里需要确保 Unix域套接字已连接  
    connect(unix_socket, (struct sockaddr*)&addr, sizeof(addr));  
    
    char message[256];  
    snprintf(message, sizeof(message), "%s:%d", client_ip, ntohs(address.sin_port));  
    send(unix_socket, message, strlen(message), 0);  
    close(unix_socket);  
    uint8_t buffer[256];
    ssize_t bytes_received;
    
    // 设置超时（10秒）
    struct timeval timeout = {10, 0};
    setsockopt(client_socket, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
    
    while ((bytes_received = recv(client_socket, buffer, sizeof(buffer), 0)) > 0) {
        // 解析请求头
        uint16_t transaction_id = (buffer[0] << 8) | buffer[1];
        uint8_t unit_id = buffer[6];
        uint8_t function_code = buffer[7];
        uint16_t starting_address = (buffer[8] << 8) | buffer[9];
        uint16_t quantity = (buffer[10] << 8) | buffer[11];

        // 准备响应头
        uint8_t response[256];
        response[0] = (transaction_id >> 8) & 0xFF;
        response[1] = transaction_id & 0xFF;
        response[2] = 0x00;
        response[3] = 0x00;
        response[6] = unit_id;
        response[7] = function_code;

        pthread_mutex_lock(&data_mutex);
        
        switch (function_code) {
            case 0x01: { // Read Coils
                uint8_t byte_count = (quantity + 7) / 8;
                response[5] = 3 + byte_count;
                response[8] = byte_count;
                
                for (int i = 0; i < byte_count; i++) {
                    response[9 + i] = 0;
                    for (int j = 0; j < 8 && (i*8 + j) < quantity; j++) {
                        if (starting_address + i*8 + j < COIL_COUNT) {
                            if (coils[starting_address + i*8 + j]) {
                                response[9 + i] |= (1 << j);
                            }
                        }
                    }
                }
                send(client_socket, response, 9 + byte_count, 0);
                break;
            }
            
            case 0x02: { // Read Discrete Inputs
                uint8_t byte_count = (quantity + 7) / 8;
                response[5] = 3 + byte_count;
                response[8] = byte_count;
                
                for (int i = 0; i < byte_count; i++) {
                    response[9 + i] = 0;
                    for (int j = 0; j < 8 && (i*8 + j) < quantity; j++) {
                        if (starting_address + i*8 + j < DISCRETE_INPUT_COUNT) {
                            if (discrete_inputs[starting_address + i*8 + j]) {
                                response[9 + i] |= (1 << j);
                            }
                        }
                    }
                }
                send(client_socket, response, 9 + byte_count, 0);
                break;
            }
            
            case 0x03: { // Read Holding Registers
                uint8_t byte_count = quantity * 2;
                response[5] = 3 + byte_count;
                response[8] = byte_count;
                
                for (int i = 0; i < quantity; i++) {
                    if (starting_address + i < HOLDING_REGISTER_COUNT) {
                        response[9 + i*2] = (holding_registers[starting_address + i] >> 8) & 0xFF;
                        response[10 + i*2] = holding_registers[starting_address + i] & 0xFF;
                    } else {
                        response[9 + i*2] = 0;
                        response[10 + i*2] = 0;
                    }
                }
                send(client_socket, response, 9 + byte_count, 0);
                break;
            }
            
            case 0x04: { // Read Input Registers
                uint8_t byte_count = quantity * 2;
                response[5] = 3 + byte_count;
                response[8] = byte_count;
                
                for (int i = 0; i < quantity; i++) {
                    if (starting_address + i < INPUT_REGISTER_COUNT) {
                        response[9 + i*2] = (input_registers[starting_address + i] >> 8) & 0xFF;
                        response[10 + i*2] = input_registers[starting_address + i] & 0xFF;
                    } else {
                        response[9 + i*2] = 0;
                        response[10 + i*2] = 0;
                    }
                }
                send(client_socket, response, 9 + byte_count, 0);
                break;
            }
            
            case 0x05: { // Write Single Coil
                uint8_t value = (buffer[10] == 0xFF) ? 1 : 0;
                if (starting_address < COIL_COUNT) {
                    coils[starting_address] = value;
                }
                send(client_socket, buffer, bytes_received, 0);
                break;
            }
            
            case 0x06: { // Write Single Register
                uint16_t value = (buffer[10] << 8) | buffer[11];
                if (starting_address < HOLDING_REGISTER_COUNT) {
                    holding_registers[starting_address] = value;
                }
                send(client_socket, buffer, bytes_received, 0);
                break;
            }
            
            case 0x0F: { // Write Multiple Coils
                uint8_t byte_count = buffer[12];
                uint16_t coil_count = (buffer[10] << 8) | buffer[11];
                
                for (int i = 0; i < coil_count; i++) {
                    if (starting_address + i < COIL_COUNT) {
                        uint8_t byte_index = i / 8;
                        uint8_t bit_index = i % 8;
                        if (byte_index < byte_count) {
                            coils[starting_address + i] = (buffer[13 + byte_index] & (1 << bit_index)) ? 1 : 0;
                        }
                    }
                }
                memcpy(response, buffer, 12);
                send(client_socket, response, 12, 0);
                break;
            }
            
            case 0x10: { // Write Multiple Registers
                uint8_t byte_count = buffer[12];
                uint16_t register_count = (buffer[10] << 8) | buffer[11];
                
                for (int i = 0; i < register_count; i++) {
                    if (starting_address + i < HOLDING_REGISTER_COUNT) {
                        holding_registers[starting_address + i] = (buffer[13 + i*2] << 8) | buffer[14 + i*2];
                    }
                }
                memcpy(response, buffer, 12);
                send(client_socket, response, 12, 0);
                break;
            }
            
            default: { // 不支持的Function Code
                response[5] = 0x03;
                response[7] |= 0x80;
                response[8] = 0x01;
                send(client_socket, response, 9, 0);
                break;
            }
        }
        
        pthread_mutex_unlock(&data_mutex);
    }
    
    close(client_socket);
    return NULL;
}

int main() {
    signal(SIGINT, sig_handler);
    int server_fd, client_socket;
    struct sockaddr_in address;
    int addrlen = sizeof(address);
    
    // 创建socket
    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) == 0) {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }
    
    // 设置socket选项
    int opt = 1;
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt))) {
        perror("setsockopt");
        exit(EXIT_FAILURE);
    }
    
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(PORT);
    
    // 绑定
    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("bind failed");
        exit(EXIT_FAILURE);
    }
    
    // 监听
    if (listen(server_fd, MAX_CLIENTS) < 0) {
        perror("listen");
        exit(EXIT_FAILURE);
    }
    
    printf("Modbus TCP Server started on port %d\n", PORT);
    printf("All data initialized to 0\n");
    printf("Waiting for connections...\n");
    
    while (1) {
        if ((client_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) {
            perror("accept");
            continue;
        }
        
        // printf("New connection from %s:%d\n", 
        //        inet_ntoa(address.sin_addr), ntohs(address.sin_port));

        // if (strcmp(inet_ntoa(address.sin_addr), "127.0.0.1") != 0) {
        //     printf("inet_ntoa(address.sin_addr):%s\n)",inet_ntoa(address.sin_addr));
        // }
        
        pthread_t thread;
        if (pthread_create(&thread, NULL, handle_client, (void *)(intptr_t)client_socket) != 0) {
            perror("pthread_create failed");
            close(client_socket);
        }
        
        pthread_detach(thread);
    }
    
    return 0;
}